// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef BASE_TRACE_EVENT_HEAP_PROFILER_ALLOCATION_REGISTER_H_
#define BASE_TRACE_EVENT_HEAP_PROFILER_ALLOCATION_REGISTER_H_

#include <stddef.h>
#include <stdint.h>

#include <utility>

#include "base/bits.h"
#include "base/logging.h"
#include "base/macros.h"
#include "base/process/process_metrics.h"
#include "base/template_util.h"
#include "base/trace_event/heap_profiler_allocation_context.h"

namespace base {
namespace trace_event {

    class AllocationRegisterTest;

    namespace internal {

        // Allocates a region of virtual address space of |size| rounded up to the
        // system page size. The memory is zeroed by the system. A guard page is
        // added after the end.
        void* AllocateGuardedVirtualMemory(size_t size);

        // Frees a region of virtual address space allocated by a call to
        // |AllocateVirtualMemory|.
        void FreeGuardedVirtualMemory(void* address, size_t allocated_size);

        // Hash map that mmaps memory only once in the constructor. Its API is
        // similar to std::unordered_map, only index (KVIndex) is used to address
        template <size_t NumBuckets, class Key, class Value, class KeyHasher>
        class FixedHashMap {
            // To keep things simple we don't call destructors.
            static_assert(is_trivially_destructible<Key>::value && is_trivially_destructible<Value>::value,
                "Key and Value shouldn't have destructors");

        public:
            using KVPair = std::pair<const Key, Value>;

            // For implementation simplicity API uses integer index instead
            // of iterators. Most operations (except FindValidIndex) on KVIndex
            // are O(1).
            using KVIndex = size_t;
            static const KVIndex kInvalidKVIndex = static_cast<KVIndex>(-1);

            // Capacity controls how many items this hash map can hold, and largely
            // affects memory footprint.
            FixedHashMap(size_t capacity)
                : num_cells_(capacity)
                , cells_(static_cast<Cell*>(
                      AllocateGuardedVirtualMemory(num_cells_ * sizeof(Cell))))
                , buckets_(static_cast<Bucket*>(
                      AllocateGuardedVirtualMemory(NumBuckets * sizeof(Bucket))))
                , free_list_(nullptr)
                , next_unused_cell_(0)
            {
            }

            ~FixedHashMap()
            {
                FreeGuardedVirtualMemory(cells_, num_cells_ * sizeof(Cell));
                FreeGuardedVirtualMemory(buckets_, NumBuckets * sizeof(Bucket));
            }

            std::pair<KVIndex, bool> Insert(const Key& key, const Value& value)
            {
                Cell** p_cell = Lookup(key);
                Cell* cell = *p_cell;
                if (cell) {
                    return { static_cast<KVIndex>(cell - cells_), false }; // not inserted
                }

                // Get a free cell and link it.
                *p_cell = cell = GetFreeCell();
                cell->p_prev = p_cell;
                cell->next = nullptr;

                // Initialize key/value pair. Since key is 'const Key' this is the
                // only way to initialize it.
                new (&cell->kv) KVPair(key, value);

                return { static_cast<KVIndex>(cell - cells_), true }; // inserted
            }

            void Remove(KVIndex index)
            {
                DCHECK_LT(index, next_unused_cell_);

                Cell* cell = &cells_[index];

                // Unlink the cell.
                *cell->p_prev = cell->next;
                if (cell->next) {
                    cell->next->p_prev = cell->p_prev;
                }
                cell->p_prev = nullptr; // mark as free

                // Add it to the free list.
                cell->next = free_list_;
                free_list_ = cell;
            }

            KVIndex Find(const Key& key) const
            {
                Cell* cell = *Lookup(key);
                return cell ? static_cast<KVIndex>(cell - cells_) : kInvalidKVIndex;
            }

            KVPair& Get(KVIndex index)
            {
                return cells_[index].kv;
            }

            const KVPair& Get(KVIndex index) const
            {
                return cells_[index].kv;
            }

            // Finds next index that has a KVPair associated with it. Search starts
            // with the specified index. Returns kInvalidKVIndex if nothing was found.
            // To find the first valid index, call this function with 0. Continue
            // calling with the last_index + 1 until kInvalidKVIndex is returned.
            KVIndex Next(KVIndex index) const
            {
                for (; index < next_unused_cell_; ++index) {
                    if (cells_[index].p_prev) {
                        return index;
                    }
                }
                return kInvalidKVIndex;
            }

            // Estimates number of bytes used in allocated memory regions.
            size_t EstimateUsedMemory() const
            {
                size_t page_size = base::GetPageSize();
                // |next_unused_cell_| is the first cell that wasn't touched, i.e.
                // it's the number of touched cells.
                return bits::Align(sizeof(Cell) * next_unused_cell_, page_size) + bits::Align(sizeof(Bucket) * NumBuckets, page_size);
            }

        private:
            friend base::trace_event::AllocationRegisterTest;

            struct Cell {
                KVPair kv;
                Cell* next;

                // Conceptually this is |prev| in a doubly linked list. However, buckets
                // also participate in the bucket's cell list - they point to the list's
                // head and also need to be linked / unlinked properly. To treat these two
                // cases uniformly, instead of |prev| we're storing "pointer to a Cell*
                // that points to this Cell" kind of thing. So |p_prev| points to a bucket
                // for the first cell in a list, and points to |next| of the previous cell
                // for any other cell. With that Lookup() is the only function that handles
                // buckets / cells differently.
                // If |p_prev| is nullptr, the cell is in the free list.
                Cell** p_prev;
            };

            using Bucket = Cell*;

            // Returns a pointer to the cell that contains or should contain the entry
            // for |key|. The pointer may point at an element of |buckets_| or at the
            // |next| member of an element of |cells_|.
            Cell** Lookup(const Key& key) const
            {
                // The list head is in |buckets_| at the hash offset.
                Cell** p_cell = &buckets_[Hash(key)];

                // Chase down the list until the cell that holds |key| is found,
                // or until the list ends.
                while (*p_cell && (*p_cell)->kv.first != key) {
                    p_cell = &(*p_cell)->next;
                }

                return p_cell;
            }

            // Returns a cell that is not being used to store an entry (either by
            // recycling from the free list or by taking a fresh cell).
            Cell* GetFreeCell()
            {
                // First try to re-use a cell from the free list.
                if (free_list_) {
                    Cell* cell = free_list_;
                    free_list_ = cell->next;
                    return cell;
                }

                // Otherwise pick the next cell that has not been touched before.
                size_t idx = next_unused_cell_;
                next_unused_cell_++;

                // If the hash table has too little capacity (when too little address space
                // was reserved for |cells_|), |next_unused_cell_| can be an index outside
                // of the allocated storage. A guard page is allocated there to crash the
                // program in that case. There are alternative solutions:
                // - Deal with it, increase capacity by reallocating |cells_|.
                // - Refuse to insert and let the caller deal with it.
                // Because free cells are re-used before accessing fresh cells with a higher
                // index, and because reserving address space without touching it is cheap,
                // the simplest solution is to just allocate a humongous chunk of address
                // space.

                DCHECK_LT(next_unused_cell_, num_cells_ + 1);

                return &cells_[idx];
            }

            // Returns a value in the range [0, NumBuckets - 1] (inclusive).
            size_t Hash(const Key& key) const
            {
                if (NumBuckets == (NumBuckets & ~(NumBuckets - 1))) {
                    // NumBuckets is a power of 2.
                    return KeyHasher()(key) & (NumBuckets - 1);
                } else {
                    return KeyHasher()(key) % NumBuckets;
                }
            }

            // Number of cells.
            size_t const num_cells_;

            // The array of cells. This array is backed by mmapped memory. Lower indices
            // are accessed first, higher indices are accessed only when the |free_list_|
            // is empty. This is to minimize the amount of resident memory used.
            Cell* const cells_;

            // The array of buckets (pointers into |cells_|). |buckets_[Hash(key)]| will
            // contain the pointer to the linked list of cells for |Hash(key)|.
            // This array is backed by mmapped memory.
            mutable Bucket* buckets_;

            // The head of the free list.
            Cell* free_list_;

            // The index of the first element of |cells_| that has not been used before.
            // If the free list is empty and a new cell is needed, the cell at this index
            // is used. This is the high water mark for the number of entries stored.
            size_t next_unused_cell_;

            DISALLOW_COPY_AND_ASSIGN(FixedHashMap);
        };

    } // namespace internal

    class TraceEventMemoryOverhead;

    // The allocation register keeps track of all allocations that have not been
    // freed. Internally it has two hashtables: one for Backtraces and one for
    // actual allocations. Sizes of both hashtables are fixed, and this class
    // allocates (mmaps) only in its constructor.
    class BASE_EXPORT AllocationRegister {
    public:
        // Details about an allocation.
        struct Allocation {
            const void* address;
            size_t size;
            AllocationContext context;
        };

        // An iterator that iterates entries in no particular order.
        class BASE_EXPORT ConstIterator {
        public:
            void operator++();
            bool operator!=(const ConstIterator& other) const;
            Allocation operator*() const;

        private:
            friend class AllocationRegister;
            using AllocationIndex = size_t;

            ConstIterator(const AllocationRegister& alloc_register,
                AllocationIndex index);

            const AllocationRegister& register_;
            AllocationIndex index_;
        };

        AllocationRegister();
        AllocationRegister(size_t allocation_capacity, size_t backtrace_capacity);

        ~AllocationRegister();

        // Inserts allocation details into the table. If the address was present
        // already, its details are updated. |address| must not be null.
        void Insert(const void* address,
            size_t size,
            const AllocationContext& context);

        // Removes the address from the table if it is present. It is ok to call this
        // with a null pointer.
        void Remove(const void* address);

        // Finds allocation for the address and fills |out_allocation|.
        bool Get(const void* address, Allocation* out_allocation) const;

        ConstIterator begin() const;
        ConstIterator end() const;

        // Estimates memory overhead including |sizeof(AllocationRegister)|.
        void EstimateTraceMemoryOverhead(TraceEventMemoryOverhead* overhead) const;

    private:
        friend AllocationRegisterTest;

        // Expect max 1.5M allocations. Number of buckets is 2^18 for optimal
        // hashing and should be changed together with AddressHasher.
        static const size_t kAllocationBuckets = 1 << 18;
        static const size_t kAllocationCapacity = 1500000;

        // Expect max 2^15 unique backtraces. Can be changed to 2^16 without
        // needing to tweak BacktraceHasher implementation.
        static const size_t kBacktraceBuckets = 1 << 15;
        static const size_t kBacktraceCapacity = kBacktraceBuckets;

        struct BacktraceHasher {
            size_t operator()(const Backtrace& backtrace) const;
        };

        using BacktraceMap = internal::FixedHashMap<
            kBacktraceBuckets,
            Backtrace,
            size_t, // Number of references to the backtrace (the key). Incremented
            // when an allocation that references the backtrace is inserted,
            // and decremented when the allocation is removed. When the
            // number drops to zero, the backtrace is removed from the map.
            BacktraceHasher>;

        struct AllocationInfo {
            size_t size;
            const char* type_name;
            BacktraceMap::KVIndex backtrace_index;
        };

        struct AddressHasher {
            size_t operator()(const void* address) const;
        };

        using AllocationMap = internal::FixedHashMap<
            kAllocationBuckets,
            const void*,
            AllocationInfo,
            AddressHasher>;

        BacktraceMap::KVIndex InsertBacktrace(const Backtrace& backtrace);
        void RemoveBacktrace(BacktraceMap::KVIndex index);

        Allocation GetAllocation(AllocationMap::KVIndex) const;

        AllocationMap allocations_;
        BacktraceMap backtraces_;

        DISALLOW_COPY_AND_ASSIGN(AllocationRegister);
    };

} // namespace trace_event
} // namespace base

#endif // BASE_TRACE_EVENT_HEAP_PROFILER_ALLOCATION_REGISTER_H_
